import pandas as pd
import numpy as np
import logging
from pathlib import Path


class OlympicMedalPredictor:
    def __init__(self):
        # 初始化日志
        logging.basicConfig(level=logging.INFO)
        self.logger = logging.getLogger(__name__)

    def enhance_athlete_features(self, data):
        """增强运动员特征"""
        features = data.copy()

        # 1. 时间衰减权重
        current_year = 2024
        features['recency_weight'] = np.exp(-0.1 * (current_year - data['last_year']))
        features['weighted_performance'] = features['avg_performance'] * features['recency_weight']

        # 2. 成绩稳定性指标
        features['performance_consistency'] = data['avg_performance'] / data['best_performance'].clip(lower=0.001)

        # 3. 参赛密度和频率
        features['participation_frequency'] = data['participation_count'] / \
                                              (data['last_year'] - data['first_year']).clip(lower=1)

        # 4. 奖牌效率
        features['medal_efficiency'] = data['total_medals'] / data['participation_count'].clip(lower=1)

        # 确保country列保留在结果中
        if 'country' in data.columns:
            features = features[['Name', 'country'] + [col for col in features.columns if col not in ['Name', 'country']]]

        return features

    def enhance_features(self):
        """执行特征增强"""
        try:
            # 读取原有特征数据
            features_path = Path('features')
            self.athlete_data = pd.read_csv(features_path / 'athlete_features.csv')

            # 增强运动员特征
            enhanced_features = self.enhance_athlete_features(self.athlete_data)

            # 保存增强后的特征
            output_path = features_path / 'enhanced_athlete_features.csv'
            enhanced_features.to_csv(output_path, index=False)

            self.logger.info(f"特征增强完成，保存至: {output_path}")

            # 打印特征统计信息
            self.logger.info("\n新增特征统计:")
            new_features = ['recency_weight', 'weighted_performance',
                            'performance_consistency', 'participation_frequency',
                            'medal_efficiency']

            for feature in new_features:
                stats = enhanced_features[feature].describe()
                self.logger.info(f"\n{feature}统计信息:\n{stats}")

            return enhanced_features

        except Exception as e:
            self.logger.error(f"特征增强失败: {str(e)}")
            raise


if __name__ == "__main__":
    # 创建预测器实例
    predictor = OlympicMedalPredictor()

    try:
        # 执行特征增强
        enhanced_features = predictor.enhance_features()
        print("特征增强成功完成！")

        # 显示新特征的相关性分析
        correlation_matrix = enhanced_features[[
            'total_medals',  # 目标变量
            'recency_weight',
            'weighted_performance',
            'performance_consistency',
            'participation_frequency',
            'medal_efficiency'
        ]].corr()

        print("\n新特征与目标变量(total_medals)的相关性:")
        print(correlation_matrix['total_medals'].sort_values(ascending=False))

    except Exception as e:
        print(f"错误: {str(e)}")
